# 假设已经进行了从e到o的坐标转换
import numpy as np
from scipy.spatial.transform import Rotation


###################################
def iP_2_oP(oR_i, oP_iORG, iP):
    oP = oR_i@iP + oP_iORG

    return oP


def triangle_angles_A(a, b, c):

    # 计算三个角的余弦值
    cos_A = (b**2 + c**2 - a**2) / (2 * b * c)
    
    # 转换为弧度
    A = np.arccos(cos_A)
    
    return A


#############################

# 这里就是直接调用IK，传参是mocap点的位置和当前关节空间位置
# 然后把结果在mujoco里面显示出来

import mujoco
import mujoco.viewer as viewer
from os.path import dirname, join
import time

# 导入自定义函数和CANFD SDK动态链接库>>>
from os.path import dirname, join
import sys
import platform
arch = platform.machine()
sys.path.append(join(dirname(__file__), 'src/lysdemo_'+arch))
sys.path.append(join(dirname(__file__), "src"))
sys.path.append(dirname(__file__))
from src import *

pos_sensor_names = [
        "Motor_Shoulder_pos", 
        "Motor_BigArm1_pos",
        "Motor_SmallArm1_pos",
        "Motor_SmallArm2_pos",
        "Motor_Hand_pos",
        "Motor_Finger_pos"
        ]

vel_sensor_names = [
        "Motor_Shoulder_vel", 
        "Motor_BigArm1_vel",
        "Motor_SmallArm1_vel",
        "Motor_SmallArm2_vel",
        "Motor_Hand_vel",
        "Motor_Finger_vel"
        ]

# mujoco加载模型
model_dof6 = mujoco.MjModel.from_xml_path(join(dirname(__file__),'scene_6DOF_new.xml')) 
data_dof6 = mujoco.MjData(model_dof6)
mocap_id_dof6 = model_dof6.body("target").mocapid[0]




#####################################
l_0 = 1
l_1 = 0.135
l_2=0.5
l_3=0.511
l_5 = 0.202

# 坐标D是目标的本体坐标系
DP_C = np.array([0, 0, -l_5])





oT_e = np.array([[0, -1, 0, 0], [1, 0, 0, 0], [0, 0, 1, -l_0-l_1], [0, 0, 0, 1]])
oR_e = oT_e[:3, :3]
oP_eORG = oT_e[:3, 3]

flag = 1

########################################

with viewer.launch_passive(model_dof6,data_dof6) as viewer:
    # Initialize the camera view to that of the free camera.
    mujoco.mjv_defaultFreeCamera(model_dof6, viewer.cam)
    viewer.opt.frame = mujoco.mjtFrame.mjFRAME_SITE


    while viewer.is_running():
        q_real_float = np.array([data_dof6.sensor(pos_sensor_name).data[0] for pos_sensor_name in pos_sensor_names])
        qv_real_float = np.array([data_dof6.sensor(vel_sensor_name).data[0] for vel_sensor_name in vel_sensor_names])

        if flag == 1:
            last_qpos = q_real_float

        e_P_DORG = data_dof6.mocap_pos[mocap_id_dof6]
        D_quat = data_dof6.mocap_quat[mocap_id_dof6]
        eR_D = Rotation.from_quat([D_quat[1], D_quat[2], D_quat[3], D_quat[0]]).as_matrix()



        # IK_ANALITICAL
        # 还要判断一下各个角度的范围，再据此调整flag
        ############################
        time_start = time.time()
        # 求解123
        eP_C = iP_2_oP(eR_D, e_P_DORG, DP_C)
        oP_C = iP_2_oP(oR_e, oP_eORG, eP_C)
        oP_AC = oP_C

        theta_1 = np.arctan2(oP_AC[1], oP_AC[0])

        P_AC_norm = np.linalg.norm(oP_AC)
        Angle_AC_Axis1 = np.arccos(oP_AC[2] / P_AC_norm)
        if P_AC_norm > (l_2 + l_3):
            print("too far")
            flag = 0 # 设置一个flag = 0，该条件下，所有关节角度不变
        elif Angle_AC_Axis1 < 0.1:
            flag = 0
            print("123singular")
        else:
            flag = 1
            theta_3 = math.pi - triangle_angles_A(P_AC_norm, l_2, l_3)
            # 求P_{AC}和Axis_1的夹角, arccos的范围是0~pi
            theta_2 = Angle_AC_Axis1-triangle_angles_A(l_3, P_AC_norm, l_2)
        


        # 求解456
        if flag == 1:
            # 计算矩阵^{o}R_{C0}
            oP_B = np.array([l_2*math.sin(theta_2)*math.cos(theta_1),l_2*math.sin(theta_2)*math.sin(theta_1), l_2*math.cos(theta_2)])
            oP_BC = oP_C-oP_B
            oZ_C0 = oP_BC/np.linalg.norm(oP_BC)
            oX_C0 = np.array([-math.sin(theta_1), math.cos(theta_1), 0])
            oY_C0 = np.cross(oZ_C0,oX_C0)
            oR_C0 = np.column_stack((oX_C0, oY_C0, oZ_C0))

            # 计算矩阵^{C0}R_{D}
            oR_D = oR_e@eR_D
            C0R_D = oR_C0.T@oR_D
            r = C0R_D
            
            # 计算theta456
            theta_5 = np.arctan2(math.sqrt(r[2][0]**2+r[2][1]**2),r[2][2])
            if theta_5 == 0:
                theta_4 = last_qpos[3]
                theta_6 = np.arctan2(-r[0][1],r[0][0]) - theta_4
            else:
                theta_4 = np.arctan2(r[1][2]/math.sin(theta_5), r[0][2]/math.sin(theta_5))
                theta_6 = np.arctan2(r[2][1]/math.sin(theta_5), -r[2][0]/math.sin(theta_5))


        print(f"time  cost: {time.time()-time_start}")
        ###################################




        if flag == 0:
            data_dof6.ctrl = last_qpos
        else:
            data_dof6.ctrl = np.array([theta_1,theta_2,theta_3,theta_4,theta_5,theta_6])
        
        mujoco.mj_step(model_dof6, data_dof6)
        mujoco.mj_forward(model_dof6, data_dof6)

        viewer.sync()
        time.sleep(model_dof6.opt.timestep)